package io.github.eric.urpc.core.thread;

import java.util.Collection;
import java.util.List;
import java.util.concurrent.*;

/**
 * 有序线程池，计算key的哈希值对线程数取余，确保同哈希值的任务始终用同一个线程处理
 */
public class OrderedThreadPoolExecutor implements ExecutorService {

    private final ExecutorService[] executorServices;
    // 允许脏读，不保证一定是轮询
    private int roundRobbinIndex;

    public OrderedThreadPoolExecutor(int nThreads, int queueSize, String threadNamePrefix) {
        assert nThreads > 0;
        executorServices = new ExecutorService[nThreads];
        for (int i = 0; i < nThreads; i++) {
            ExecutorService executorService = new ThreadPoolExecutor(1, 1,
                    0L, TimeUnit.MILLISECONDS,
                    new LinkedBlockingQueue<>(queueSize),
                    new NamePrefixThreadFactory(threadNamePrefix + "-ordered-" + i));
            executorServices[i] = executorService;
        }
    }

    @Override
    public void shutdown() {
        for (ExecutorService executorService : executorServices) {
            executorService.shutdown();
        }
    }

    @Override
    public List<Runnable> shutdownNow() {
        throw new UnsupportedOperationException();
    }

    @Override
    public boolean isShutdown() {
        throw new UnsupportedOperationException();
    }

    @Override
    public boolean isTerminated() {
        throw new UnsupportedOperationException();
    }

    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
        throw new UnsupportedOperationException();
    }

    @Override
    public <T> Future<T> submit(Callable<T> task) {
        throw new UnsupportedOperationException();
    }

    @Override
    public <T> Future<T> submit(Runnable task, T result) {
        throw new UnsupportedOperationException();
    }

    @Override
    public Future<?> submit(Runnable task) {
        throw new UnsupportedOperationException();
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks) throws InterruptedException {
        throw new UnsupportedOperationException();
    }

    @Override
    public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException {
        throw new UnsupportedOperationException();
    }

    @Override
    public <T> T invokeAny(Collection<? extends Callable<T>> tasks) throws InterruptedException, ExecutionException {
        throw new UnsupportedOperationException();
    }

    @Override
    public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
        throw new UnsupportedOperationException();
    }

    @Override
    public void execute(Runnable command) {
        int index;
        String orderedKey = OrderedThreadKeyContext.getCurrent();
        if (orderedKey == null) {
            index = roundRobbinIndex++ % executorServices.length;
            if (roundRobbinIndex == Integer.MAX_VALUE/2) {
                roundRobbinIndex = 0;
            }
        } else {
            index = orderedKey.hashCode() % executorServices.length;
        }
        executorServices[index].execute(command);
    }

    public int getQueueSize() {
        int size = 0;
        for (ExecutorService executorService : executorServices) {
            size += ((ThreadPoolExecutor) executorService).getQueue().size();
        }
        return size;
    }
}
